import cv2
import sys
from tracker import *

def test_one_img(video_path):
  # 创建跟踪器对象
  img = cv2.imread(video_path)
  # 来自稳定相机的物体检测
  object_detector = cv2.createBackgroundSubtractorMOG2(history=10, varThreshold=10.)
  roi = img

  # 1. 物体检测
  mask = object_detector.apply(roi)
  # _, mask = cv2.threshold(mask, 254, 255, cv2.THRESH_BINARY)
  contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
  detections = []
  for cnt in contours:
    # 计算面积并去除小元素
    area = cv2.contourArea(cnt)
    if area > 100:
      #cv2.drawContours(roi, [cnt], -1, (0, 255, 0), 2)
      x, y, w, h = cv2.boundingRect(cnt)
      detections.append([x, y, w, h])

  # 2. 对象追踪
  for box in detections:
    x, y, w, h = box
    cv2.rectangle(roi, (x, y), (x + w, y + h), (0, 255, 0), 3)
  cv2.imwrite("roi.jpg", roi)
  cv2.imwrite("mask.jpg", mask)
  # cv2.imshow("Frame", frame)
  # cv2.imshow("Mask", mask)

  

def main(video_path):
  # 创建跟踪器对象
  tracker = EuclideanDistTracker()
  cap = cv2.VideoCapture(video_path)
  # 来自稳定相机的物体检测
  object_detector = cv2.createBackgroundSubtractorMOG2(history=100, varThreshold=40)
  while True:
    ret, frame = cap.read()
    height, width, _ = frame.shape
    # 提取感兴趣区域
    # roi = frame[340: 720,500: 800]
    roi = frame

    # 1. 物体检测
    mask = object_detector.apply(roi)
    _, mask = cv2.threshold(mask, 254, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    detections = []
    for cnt in contours:
      # 计算面积并去除小元素
      area = cv2.contourArea(cnt)
      if area > 100:
        #cv2.drawContours(roi, [cnt], -1, (0, 255, 0), 2)
        x, y, w, h = cv2.boundingRect(cnt)
        detections.append([x, y, w, h])
    # 2. 对象追踪
    boxes_ids = tracker.update(detections)
    for box_id in boxes_ids:
      x, y, w, h, id = box_id
      cv2.putText(roi, str(id), (x, y - 15), cv2.FONT_HERSHEY_PLAIN, 2, (255, 0, 0), 2)
      cv2.rectangle(roi, (x, y), (x + w, y + h), (0, 255, 0), 3)
    cv2.imshow("roi", roi)
    # cv2.imshow("Frame", frame)
    # cv2.imshow("Mask", mask)
    key = cv2.waitKey(30)
    if key == 27:
      break
  cap.release()
  cv2.destroyAllWindows()


if __name__ == "__main__":
  video_path = sys.argv[1]
  test_one_img(video_path)
